"""
利用radare2进行数据处理，获得functionSim的模型输入
"""
import json
import signal
import r2pipe
import subprocess
import datetime
import shelve
from tqdm import tqdm
from bin2asm import *
import os
import sys
sys.path.append(
    r"/home/cyw/projects/function_sim_project/siamese_Graphsage/asm2vec-pytorch-master/scripts")


def set_timeout(num, callback):
    """
        函数超时设置，避免一直卡在一个函数中
    """
    def wrap(func):
        # 收到信号 SIGALRM 后的回调函数，第一个参数是信号的数字，第二个参数是the interrupted stack frame.
        def handle(signum, frame):
            raise RuntimeError

        def to_do(*args, **kwargs):
            try:
                signal.signal(signal.SIGALRM, handle)  # 设置信号和回调函数
                signal.alarm(num)  # 设置 num 秒的闹钟
                print('start alarm signal.')
                r = func(*args, **kwargs)
                print('close alarm signal.')
                signal.alarm(0)  # 关闭闹钟
                return r
            except RuntimeError as e:
                callback()
        return to_do
    return wrap


def after_timeout():  # 超时后的处理函数
    print("Time out!")


timeOut = 1600

logPath = r"/home/cyw/projects/function_sim_project/disassemblyTools/radare2/radare2_logs.txt"


def logs(inf):
    with open(logPath, "a+") as file:
        file.write(str(datetime.datetime.now())+"\t" + inf + "\n")


def getUnfinishedName(originPath, savePath):
    sampleList = {}
    finishList = {}
    unfinList = {}
    failedList = {}
    sLoc = 0
    fLoc = 0
    uLoc = 0
    # 获取所有的样本名
    for file in os.listdir(originPath):
        data = file.split(".")  # 找到样本文件
        if (len(data) == 1):
            sLoc += 1
            sampleList[data[0]] = True
    # 获取已经处理过的样本名
    for file in os.listdir(savePath):
        data = file.split(".")
        if (data[1] == "bak"):
            fLoc += 1
            finishList[data[0]] = True
    # 获得不能处理的样本
    path = "/home/cyw/projects/function_sim_project/disassemblyTools/radare2/failedSample.txt"
    with open(path) as file:
        for str in file.readlines():
            failedList[str.strip()] = True
    # 获得未被处理过的样本
    for name in sampleList.keys():
        if name not in finishList and name not in failedList:
            uLoc += 1
            unfinList[name] = True
    logs("样本总计{}个，已处理{}个,还剩{}个\n".format(sLoc, fLoc, uLoc))
    return unfinList


class radare2DevTools:
    dataTrans = ["mov", "push", "pop", "xchg", "in", "out", "xlat",
                 "cwd", "lea", "lds", "les", "lahf", "sahf", "pushf", "popf"]
    arithmetric = ["add", "adc", "inc", "daa", "sub", "sbb", "dec", "cmp",
                   "das", "mul", "imul", "div", "idiv", "xor", "or", "not", "sll", "srl"]
    controlTrans = ["call", "ret", "jmp", "loop", "loope", "loopne", "int", "into", "iret", "jz",
                    "je", "jnz", "jne", "jo", "jno", "jc", "jnc", "js", "jns", "jb", "jnae", "jnb", "jae", "ja",
                    "jnbe", "jna", "jbe", "jl", "jnge", "jnl", "jge", "jg", "jnle", "jng", "jle"]
    success = False
    DYNAMICFUNCTION = "dynamic import"
    STATICFUNCTION = "static link"
    LOCALFUNCTION = "local"

    def __init__(self) -> None:
        pass

    def save_data(self, savePath, res):
        # 保存数据
        file = shelve.open(savePath)
        file["cg"] = res["cg"]
        file["cgattr"] = res["cgattr"]
        file["functype"] = res["functype"]
        file["funcs_id"] = res["funcs_id"]
        file.close()
        print("样本保存成功")

    def gene_cg(self, functions, funcs_id, r2):
        """
            获得函数的邻接矩阵
        """
        calls = {}
        funcSize = 0
        for function in functions:
            name = function["name"]
            calls[name] = []
            funcSize += 1

        cg = []
        for i in range(funcSize):
            cg.append([0]*funcSize)
        # 获得函数之间的调用关系
        for function in functions:
            name = function["name"]
            address = function["offset"]
            # axf 该函数引用的所有地方(什么结果都没有，奇怪)  axt引用了对应地址的地方
            xrefs = r2.cmdj(f"axtj {address}")
            for xref in xrefs:
                if xref["type"] == "CALL":
                    # 加@表示地址，否则全为空
                    called_function_name = r2.cmd(
                        f"afn @{xref['from']}").strip()
                    # 特例处理
                    if called_function_name == "":
                        continue
                    # 这里反着写，逻辑才对  calls[name].append(called_function_name)
                    # 新问题修复,2个样本会出现keyError的情况,这里做个特例处理
                    if called_function_name in calls:
                        calls[called_function_name].append(name)
        # 生成邻接矩阵
        for name in calls.keys():
            ind = funcs_id[name]
            for subName in calls[name]:
                if subName == "":
                    continue
                indy = funcs_id[subName]
                cg[ind][indy] += 1
        return cg

    def gene_asm_code(self, r2):
        """
            传输r2的pdf信息，从中提取汇编代码
            优化处理
        """
        tar = r2.cmd('pdf')
        asms = tar.split("\n")
        lth = len(asms)
        lims = ["; XREFS", "; CALL XREF",
                "; CODE XREF", "; var", ";-- eip:", "; arg"]
        # :前为数字
        # 艰难的匹配过程
        res = []
        for i in range(lth):
            tar = asms[i]
            if tar == "" or len(tar) <= 5:
                continue
            is_asm = True
            for lim in lims:
                if lim in tar:
                    is_asm = False
            tempNum = tar.split(":")[0].split(" ")[1]
            if (tempNum.isdigit()):
                is_asm = False
            if (is_asm):
                res.append(tar[43:])

        return res

    def gene_attr(self, r2, functions, cg, funcType):
        """
            获得每个函数的嵌入向量
            # 0：字符常量数量
            # 1：数值常量数量
            # 2：转移指令数量
            # 3：控制指令数量
            # 4：指令数量
            # 5：算术指令数量
            # 6：后代数量
            # 7：上一辈数量
        """
        cgattr = []
        # 还需要获得每个函数的汇编代码
        for ind in range(len(functions)):
            function = functions[ind]
            temp = [0, 0, 0, 0, 0, 0, 0, 0]
            r2.cmd(f's {function["offset"]}')
            asms = self.gene_asm_code(r2)
            for asm in asms:
                tokens = asm.split(";")[0].split(";")[0].strip().split(" ")
                temp[4] += 1
                instrName = tokens[0]
                if instrName in self.dataTrans:
                    temp[2] += 1
                elif instrName in self.arithmetric:
                    temp[5] += 1
                elif instrName in self.controlTrans:
                    temp[3] += 1
                if '"' in asm or "'" in asm:
                    temp[0] += 1
                for i in range(1, len(tokens)):
                    if tokens[i].isdigit():
                        temp[1] += 1
            size = len(cg)
            for i in range(size):
                if cg[ind][i] != 0:
                    temp[6] += 1
                if cg[i][ind] != 0:
                    temp[7] += 1
            if funcType[ind] == self.DYNAMICFUNCTION:
                temp = [0, 0, 0, 0, 0, 0, 0, 0]
            cgattr.append(temp)
        return cgattr

    def getFuntionType(self, r2, functions):
        # 动态导入函数判定：函数名有dll,样本长度为1的thunk函数为动态导入函数
        # name:sub.MSVCRT.dll_memset      length:1        type:fcn
        # ['jmp dword [sym.imp.MSVCRT.dll_memset] ; 0x406174 ; "\\"\\x8a"']
        # 静态链接库函数判定

        res = {}
        for cnt in range(len(functions)):
            func = functions[cnt]
            funcType = self.STATICFUNCTION
            name = func['name']
            r2.cmd(f's {func["offset"]}')
            asms = self.gene_asm_code(r2)
            # print("name:{}\tlength:{}\ttype:{}".format(name,len(asms),func["type"]))
            # print(asms[:4])
            # print()
            if ".dll" in name and len(asms) <= 2:
                funcType = self.DYNAMICFUNCTION
            elif "fcn." in name:
                funcType = self.LOCALFUNCTION
            res[cnt] = funcType
        return res

    @set_timeout(timeOut, after_timeout)
    def analysisBinary(self, originPath, name):
        """
            进行样本的分析
            保存如下数据：
            # cg
            # cgattr
            # functypes
            # funcs_id
        """
        try:
            samplePath = originPath + "/" + name
            # Use -2 flag for non-interactive mode
            r2 = r2pipe.open(samplePath, flags=['-2'])
            r2.cmd('aaa')  # a越多分析的东西越多，asm2vec-pytorch用四个a，但是好像最多就三个
            functions = r2.cmdj('aflj')  # 获得函数列表
        except Exception as e:
            logs("异常：{}".format(e))
            self.success = False

        # 函数间的调用
        functype = {}
        cnt = 0
        funcs_id = {}
        for function in functions:
            name = function["name"]
            funcs_id[name] = cnt
            cnt += 1
        funcSize = cnt
        cg = self.gene_cg(functions, funcs_id, r2)
        functype = self.getFuntionType(r2, functions)
        cgattr = self.gene_attr(r2, functions, cg, functype)

        res = {}
        res["cgattr"] = cgattr
        res["funcs_id"] = funcs_id
        res["cg"] = cg
        res["functype"] = functype
        return res

    # 分析文件夹下的样本，
    def analysisPath(self, originPath, savePath):
        unfinishedList = getUnfinishedName(originPath, savePath)
        logs("使用radare2,进行样本反汇编")
        for name in tqdm(unfinishedList):
            logs("样本{}正在处理".format(name))
            self.success = True
            res = self.analysisBinary(originPath, name)
            if (self.success):
                self.save_data(savePath+"/"+name, res)
            else:
                path = "/home/cyw/projects/function_sim_project/disassemblyTools/radare2/failedSample.txt"
                with open(path, "a+") as file:
                    file.write(name+"\n")
                logs("radare2处理失败:{}".format(name))
        logs("样本反汇编完成")


if __name__ == '__main__':
    originPath = r"/home/cyw/projects/data/functionsim/allorigindata"
    # originPath = r"/home/cyw/projects/data/functionsim/test"
    savePath = r"/home/cyw/projects/function_sim_project/all_data/sampleDatas/radare2_functionSim"
    radare2Tool = radare2DevTools()
    radare2Tool.analysisPath(originPath, savePath)
